package middlewares

import (
	"longmen/server/config/global"
	"sync"

	"github.com/labstack/echo/v4"
	"golang.org/x/time/rate"
)

type IPRateLimiter struct {
	ips map[string]*rate.Limiter
	mu  *sync.RWMutex
	r   rate.Limit
	b   int
}

var myMap sync.Map

var limiter = NewIPRateLimiter(5, 5)

func Limiter(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) error {
		return next(c)
		if limiter.getIp(c.RealIP()) {
			return next(c)
		}

		b, _ := global.BloomExists(global.BLOOM_IPBLACK, c.RealIP())
		if b {
			global.Log.Errorf("Limiter ip:%v", c.RealIP())
			return c.JSON(403, "IP Request Limit[您的请求过于频繁]")
		}
		protector := limiter.GetLimiter(c.RealIP())
		if !protector.Allow() {
			global.BloomAdd(global.BLOOM_IPBLACK, c.RealIP())
			global.Log.Errorf("Limiter ip:%v", c.RealIP())
			return c.JSON(403, "IP Request Limit[您的请求过于频繁]")
		}

		return next(c)
	}
}

// NewIPRateLimiter .
func NewIPRateLimiter(r rate.Limit, b int) *IPRateLimiter {
	i := &IPRateLimiter{
		ips: make(map[string]*rate.Limiter),
		mu:  &sync.RWMutex{},
		r:   r,
		b:   b,
	}
	i.setIp()
	return i
}

func (i *IPRateLimiter) setIp() {
	// allow list
	ips := []string{"222.127.7.210", "222.127.7.214", "218.189.125.19", "127.0.0.1"}
	for _, ip := range ips {
		myMap.Store(ip, 1)
	}
}

func (i *IPRateLimiter) getIp(ip string) bool {
	_, b := myMap.Load(ip)
	return b
}

// AddIP creates a new rate limiter and adds it to the ips map,
// using the IP address as the key
func (i *IPRateLimiter) AddIP(ip string) *rate.Limiter {
	i.mu.Lock()
	defer i.mu.Unlock()

	protector := rate.NewLimiter(i.r, i.b)

	i.ips[ip] = protector

	return protector
}

// GetLimiter returns the rate limiter for the provided IP address if it exists.
// Otherwise calls AddIP to add IP address to the map
func (i *IPRateLimiter) GetLimiter(ip string) *rate.Limiter {
	i.mu.Lock()
	protector, exists := i.ips[ip]

	if !exists {
		i.mu.Unlock()
		return i.AddIP(ip)
	}

	i.mu.Unlock()

	return protector
}
